{-# OPTIONS_GHC -cpp #-}
{-# LANGUAGE ScopedTypeVariables,BangPatterns #-}
module Learner where

import Simulator
import Types
import Random
import Management
import ManagementStore
import Constants

import Debug.Trace

import Statistics.RandomVariate as RV

import System.CPUTime
import System.IO
import System.IO.Error
import System.Exit
import System.FilePath

import qualified Data.ByteString.Lazy as B
import qualified Data.Map as Map
import Data.Array.Vector
import Data.Binary
import Data.List (sortBy,intersperse,intercalate)
import Data.Maybe
import Numeric

import Control.Concurrent
import Control.Concurrent.STM
import Control.Monad
import Control.Monad.ST
import Control.DeepSeq

-- take a world, randomly assign management to the banks,
-- run the simulation a bit, pick the best management (by total returns),
-- mutate the management and iterate for a bit. 

getAllMgt :: World -> [(Cash,Management)]
getAllMgt w
    = (zip (map bankScore bbanks) (map management bbanks))
    where bbanks = (sortBy (\a1 a2 -> compare (bankScore a2) (bankScore a1)) (banks w))

getBestBanks :: Int -> World -> [Bank]
getBestBanks i w
    = take i (sortBy (\a1 a2 -> compare (bankScore a2) (bankScore a1))
              (banks w))

bankScore b 
    | (regulatorControl b) = failPenalty b * failTime b
    | otherwise = (sumDividends b) -- (bookEquity b) * 0.0001 -- slight penalty for holding equity -- analysts whingeing about lazy balance sheets or something

assignManagement :: Gen a -> [Management] -> World -> ST a World
assignManagement g mgt w@(W {banks=bs})
    = do ms <- (shuffle g mgt)
         return w { banks = [ (b {management = m}) | (b,m) <- zip bs (cycle ms) ] }
      where message = "BUILDING, managers: " ++ (show$length mgt) 
                      ++ " banks: " ++ (show$length bs) ++ "\n"

runWorldSim :: Gen a -> Bool -> SimState -> ST a SimState
runWorldSim g ka s = 
    do fw <- nStepWorld g ka (worldsteps s) (world s)
       return s { world = fw }

--
--  Mutation and crossover in here (should be moved to its own module)
--
--  we mutate the old good mgt, put them through crossover, and combine the
--

resetWorldSim :: Gen a -> SimState -> SimState -> ST a SimState
resetWorldSim g starter s
    = do mutmgt1 <- mapM (mutate g (mutator s)) bestmgt
         mutmgt2 <- crossOver g (mutmgt1 ++ bestmgt)
         assmgtw <- (assignManagement g (newmgt mutmgt2) (world starter)) 
         return starter { steps = (steps s) - 1,
                          world = assmgtw, -- reset the world with different management
                          simid = (simid s),
                          mgt = newms } 
    where newms = addList (mgt s) allsimmgt
          allsimmgt = (getAllMgt (world s))  
          gennumber = (length$banks$world s) -- number of new managers to generate
          bestmgt = getBest (newms) (bound s)
          newmgt mm = take gennumber (concat [[a,b] | (a,b) <- zip bestmgt mm])

cpname o n = (d $ optOutput o) `combine` "checkpoint-"++(show n)++".lsc"
ycname o n = (d $ optOutput o) `combine` "yieldcurve-"++(show n)++".lyc"
diagname o n = (d $ optOutput o) `combine` "diagnostic-"++(show n)++".lyc"
outname o = (d $ optOutput o) `combine` "out"

d Nothing = "./"
d (Just x) = x

--
-- parallel version ... branches -> 
--

-- unfinished ... and a sketch for now


mergeWorlds :: [SimState] -> SimState
mergeWorlds = head 

-- Checkpoint frequency; Cleanup Frequency
-- ycf == Yield Curve Frequency

manyWorlds :: [Seed] -> Options -> Int -> Int -> Int -> SimState -> [SimState] -> IO [SimState]
manyWorlds g o n cpf cuf st@(Sim { steps=max }) s -- XXX stopat and max is broken needs rewrite
    = do oh <- openFile (outname o) WriteMode
         pManyWorlds g o oh n mmax cpf cuf outCrossFrequency st s
    where mmax = if isJust stopat
                 then fromJust stopat
                 else max
          stopat = optStopAt o 

pManyWorlds :: [Seed] -> Options -> Handle -> Int -> Int -> Int -> Int -> Int -> SimState -> [SimState] -> IO [SimState]
pManyWorlds !seed o oh n max cpf clf ocf starter (!ss)  
	| n > max = do outdata oh ("Done.\n" ++ (concat $ intersperse "\n" $ map finalSummary ss)) 
                       outdata oh (accSummary 1 (world$ head ss)) -- need more here ...
                       return ss
	| otherwise 
           = do (if (n `mod` cpf == 0)
                 then writeCheckPoint (cpname o n) seed n starter ss 
                 else return ())
                doyieldcurve n o ss
                starttime <- getCPUTime
                res <- runPar2 $ map (\x -> return (runST (sim1 x))) (zip3 seed ss (enumFromTo 1 (length ss)))
                let (!mnewss,!seed',!prints) = unzip3 res 
                let finss = outcrossmgt $ map (\newss -> newss { mgt = cle n clf (mgt newss)}) mnewss
                outdata oh $ concat $ intersperse "\n" $ sortBy compare $ prints
                endtime <- getCPUTime
                outdata oh ("# step " ++ (show n) ++ " time - " ++ (showf $ (fromIntegral (endtime - starttime))/10^12 ))
                pManyWorlds seed' o oh (n + 1) max cpf clf ocf starter finss
	where sim1 (seed,s',id)
	         = do g <- {-# SCC "sim1-1" #-} restore seed
                      postreset <- {-# SCC "sim1-2" #-} (resetWorldSim g starter s')
                      ss' <- {-# SCC "sim1-3" #-} runWorldSim g ka postreset
                      seed' <- {-# SCC "sim1-4" #-} save g
                      let ps = {-# SCC "sim1-5" #-} bsummary n ss'
                      ss' `seq` seed' `seq` ps `seq` return (ss',seed',ps)
              outdata :: Handle -> String -> IO ()
              outdata h s = putStrLn s >> hPutStrLn h s >> hFlush h
              doyieldcurve n' o' ss' 
                  = let ycf' = optYieldCurves o'
                        maxb = 3
                    in if isJust ycf' && n' `mod` (fromJust ycf') == 0
                       then do hc <- openFile (ycname o' n') WriteMode
                               sequence [hPutStrLn hc $ yieldCurveDiagnostic i b ((getBestBanks maxb $ world s') !! b) |
                                         (s',i) <- (zip ss' [1..length ss']), b <-[0..(maxb - 1)] ]
                               hClose hc
#ifdef DIAGNOSTICS
                               hc <- openFile (diagname o' n') WriteMode
                               mapM (\s -> hPutStrLn hc $ (wid s) ++ (finalSummary) s  ) ss'
                               mapM (\s -> hPutStrLn hc $ (wid s) ++ (diagnosticSummary . world) s ) ss'
                               hClose hc
#endif
                       else return ()
              cle n cuf ss 
                  = if n `mod` cuf == 0
                    then cleanup ss
                    else ss
              outcrossmgt sses
                  = if n `mod` ocf == 0
                    then let mstores = outCrossManagement $ map mgt sses
                             mgts ms = getBest ms (length $ banks $ world $ starter)
                         in [ starter { mgt = resetScores mstore, 
                                        world = (world starter) { banks = [ (b {management = m}) | 
                                                                            (b,m) <- zip (banks$world starter) (cycle$ mgts mstore) ] },
                                        simid = id } |
                              (mstore,id) <- zip mstores [1..] ]
                    else sses
              ka = n == max || all keepAccounts ss || 
                   n `mod` accountingFrequency == 0 -- keep accounts when n == max


wid s = "\n\n\n****** WORLD " ++ (show $ simid s) ++ " ******\n\n\n"

modifyTVar :: TVar a -> (a -> a) -> STM () 
modifyTVar tv f = readTVar tv >>= writeTVar tv . f

runPar :: [IO a] -> IO [a]
runPar ps
  = do resVar <- newTVarIO []
       mapM_ (forkIO . (>>= atomically . modifyTVar resVar . (:))) ps
       atomically $ do res <- readTVar resVar
                       when (length res < length ps) retry
                       return res

runPar2 :: [IO (SimState,Seed,String)] -> IO [(SimState,Seed,String)]
runPar2 jobs
  = do -- putStrLn "# Using runPar2 ..."
       resVar <- newTVarIO []
       mapM_ (forkIO . (runjob resVar)) jobs
       atomically $ do res <- readTVar resVar
                       when (length res < length jobs) retry
                       return res
  where runjob tv job = do !jobresult <- job
                           jobresult `seq` atomically $ do rs <- readTVar tv
                                                           writeTVar tv (jobresult : rs)

-- Turns off parallelism entirely for debugging ...
runParTest :: (NFData a) => [IO a] -> IO [a]
runParTest [job] = {-# SCC "runPartest" #-}
    do res <- job
       res `deepseq` return [res]

mean x = (sum x)/(fromIntegral (length x))

instances x [] = 0
instances x (y:ys)
	| x==y = 1+(instances x ys)
	| otherwise = instances x ys

summary :: World -> String
summary w 
    = let bestbanks = (getBestBanks 1 w)
          avgdivs = mean (map sumDividends (banks w)) 
          insolvencies = instances True (map regulatorControl (banks w))
      in " Banks: " ++ (show $ length $ banks w) ++
             "  Insolvs: " ++ (show insolvencies) ++
             "  Avg Divs: " ++ (showf avgdivs) -- ++
             -- "\nBest banks:\n" ++ (show bestbanks) 

-- a probe bond
sBond = Bond {start = zeroDate,
              maturity = (D (4,0)),
              rate = 0.0,
              principal = (10.0),
              rdefault = 0.01,
              recovery = 0.5 }

yieldCurveDiagnostic :: Int -> Int -> Bank -> String
yieldCurveDiagnostic i1 i2  b
    = (concat $ intersperse "\n" $ map (summ.checkOffLoan) offers) ++ "\n\n" ++
      (concat $ intersperse "\n" $ map (summ.checkOffDeposit) offers)
      where summ ((start,maturity,rate,principal,rdefault,recovery),r,party) = 
                (show i1) ++ " " ++ (show i2) ++ 
                (showf $ dateFloat $ maturity) ++ " " ++
                          (showf $ rate) ++ " " ++ (show r) ++ " " ++ (show party)
            offers = [ invBuildContract $ sBond { rate = r, maturity = (D (y,m)) } | 
                       r <- (enumFromThenTo 0.0 0.0025 0.1) , m <- [0..11] , y <- [0..4] ]
            -- XXX -*- XXX -*- Date == (D 20 0)
            checkOffLoan c = (c,fst $ evaluateContractP (D (20,0)) b (loanPolicy b) c cashRate 0 maxEvaluationDepth,Borrower)
            checkOffDeposit c = (c,fst $ evaluateContractP (D (20,0)) b (depPolicy b) c cashRate 0 maxEvaluationDepth,Depositor)
            cashRate = 0.04


{- briefsummary -}
bsummary :: Int -> SimState -> String
bsummary n ss 
    = let w = world ss
          id = simid ss
          bestbanks = (getBestBanks 1 w)
          allbanks = banks w
          newbanks = [ b | b <- allbanks, isNothing (lookupStatistics (mgt ss) (management b)) ]
          solvbanks = take 4 $ sortBy (\y x -> compare (sumDividends x) (sumDividends y))  $ filter (not . regulatorControl) allbanks
          avggdivs = mean $ map sumDividends $  solvbanks
          betternew = [ b |  b <- newbanks, (sumDividends b) > avggdivs ]
          avgfail = (mean (map failTime allbanks))
          avgnim = 100 * mean (map (arithmean nim) solvbanks)
          avgroe = 100 * mean (map (arithmean roe) solvbanks)
          avgroa = 100 * mean (map (arithmean roa) solvbanks)
          avgequity = mean (map bookEquity solvbanks)
          avgcap = 100 * (mean (map bookEquity solvbanks))/(mean (map riskWeightedAssets solvbanks))
          insolvencies = instances True (map regulatorControl (banks w))
          insolvnew =  instances True (map regulatorControl (newbanks))
      in intercalate " " [(show n),(show id),
                          (show $ length $ banks w),(show insolvencies),
                          (show $ length newbanks),(show insolvnew),(show $ length betternew),
                          (showf avggdivs),(showf avgnim),(showf avgroe),(showf avgroa),
                          (showf avgfail),(showf avgequity),(showf avgcap)]
             -- "\nBest banks:\n" ++ (show bestbanks) 

accountFeature :: String -> Bank -> [Double]
accountFeature s b =
    [ lup d | d <- sortBy compare (aDates a) ]
    where lup d = Map.findWithDefault 0 (d,(aid s)) a
          a = accounts b

nim = accountFeature "Net Interest Margin"
roe = accountFeature "Return on Equity"
roa = accountFeature "Return on Assets"


{- give a detailed summary of n banks -}
accSummary i w
    = let bestbanks = (getBestBanks i w)
          printB b  = (printAccs . accounts) b ++ "\n\n*** MANAGEMENT ***\n\n" ++ (show . management) b
      in concat $ map printB bestbanks 
      


geomean :: (Floating a) => (Bank -> [a]) -> Bank -> a
geomean lkup b = gmean (lkup b) 
    where gmean x = (product $ map ((+) 1) x)**(1/(fromIntegral $ length x)) - 1

arithmean :: (Floating a) => (Bank -> [a]) -> Bank -> a
arithmean lkup b = mean (lkup b) 

finalSummary :: SimState -> String
finalSummary ss
    = let bestbanks = (getBestBanks 10 (world ss))
          bestmgt = take 10 (sctoid $ mgt ss)
          stats = map ((flip Map.lookup) (idtost$mgt ss)) (map snd bestmgt)
	  diagw = (world ss)
      in "Best banks:\n" ++ (show bestbanks) ++
          "Best mgt scores:\n" ++ (mgtPrinter (zip bestmgt (map fromJust stats)))

mgtPrinter :: [((Score,MgtID),Statistics)] -> String
mgtPrinter [] = ""
mgtPrinter (((s,id),(_,_,_,k,i)):r )
           = "Score: " ++ (showf s) ++ " ID: " ++ (show id) ++ " attempts: " 
             ++ (show k) ++ " insolvs: " ++ (show i) ++ "\n" ++ (mgtPrinter r )

-- CheckPoint ...
-- Data format is:
--  "LoanSimCheckpoint", typesVersion, rngState, step, SimState

writeCheckPoint :: FilePath -> [Seed] -> Int -> SimState -> [SimState] -> IO ()
writeCheckPoint file g i starter state =
    do --putStr "Writing Checkpoint ..."
       B.writeFile file (encode ( checkPointString,typesVersion,
                            g, 
                            i,
                            starter,
                            state))
       --putStrLn " Done"

readCheckPoint :: FilePath -> IO ([Seed],Int,SimState,[SimState])
readCheckPoint file  =
    do putStr "# Reading Checkpoint ... "
       cont <- B.readFile file
       putStr ("Length is " ++ (show$ B.length cont))
       let (cps,tv,g ,i ,st,ss) = decode cont
       if (cps /= checkPointString || tv /= typesVersion)
        then ioError (userError "Corrupted Checkpoint or wrong checkpoint version")
        else do putStrLn " ... Done"
                return (g,i,st,ss)


instance NFData Seed where
  rnf x = rnf (show x) -- hack
